package com.lark.algorithm.study.segment;

/**
 * @author btmood
 * @version 1.0
 * @apiNote 线段树
 * @since 2024-06-13 17:03
 */
public class SegmentTree<E> {

    // 包含所有值元素的数组
    private E[] data;
    // 用数组表示的线段树
    private E[] tree;
    private Merger<E> merger;


    public SegmentTree(E[] arr, Merger<E> merger) {
        this.merger = merger;

        data = (E[]) new Object[arr.length];
        for (int i = 0; i < arr.length; i++) {
            data[i] = arr[i];
        }

        // 线段树的数组容量是4倍的区间容量
        tree = (E[]) new Object[4 * arr.length];

        buildSegmentTree(0, 0, data.length - 1);
    }

    /**
     * 在treeIndex的位置创建表示区间[l...r]的线段树
     * @param treeIndex 当前线段树根节点的索引
     * @param l 区间的左端点
     * @param r 区间的右端点
     */
    private void buildSegmentTree(int treeIndex, int l, int r) {
        // 当l==r，说明只有一个元素，直接让根节点等于它
        if (l == r) {
            tree[treeIndex] = data[l];
            return;
        }
        // 左孩子节点
        int leftTreeIndex = leftChild(treeIndex);
        // 右孩子节点
        int rightTreeIndex = rightChild(treeIndex);

        // 找出区间的中间索引
        // 防止出现整形溢出的情况：左边界加上距离除以2的位置
        // int mid = (l + r) / 2;
        int mid = l + (r - l) / 2;
        // 构建左子树
        buildSegmentTree(leftTreeIndex, l, mid);
        // 构建右子树
        buildSegmentTree(rightTreeIndex, mid + 1, r);

        // 这里可以通过对左右孩子节点处理，得出父节点的值；比如是孩子节点的和或者最大值、最小值等
        // 这里相当于是用了一个函数时接口
        tree[treeIndex] = merger.merge(tree[leftTreeIndex], tree[rightTreeIndex]);
    }


    /**
     * 获取容量
     * @return
     */
    public int getSize() {
        return data.length;
    }

    /**
     * 获取索引位置的元素
     * @param index
     * @return
     */
    public E get(int index) {
        if (index < 0 || index > data.length) {
            throw new IllegalArgumentException("Index is illegal");
        }
        return data[index];
    }

    /**
     * 返回完全二叉树的数组表示中，一个索引所表示的左孩子节点的索引
     * 虽然线段树本身不是完全二叉树，但是我们用数组来维护的线段树是一颗完全二叉树，空的节点表示成空。
     * @param index 索引
     * @return 左孩子节点索引
     */
    private int leftChild(int index) {
        return 2 * index + 1;
    }

    /**
     * 返回完全二叉树的数组表示中，一个索引所表示的右孩子节点的索引
     * @param index 索引
     * @return 右孩子节点索引
     */
    private int rightChild(int index) {
        return 2 * index + 2;
    }

    /**
     * 返回区间[queryL, queryR]的值
     * @param queryL 区间左边界
     * @param queryR 区间右边界
     * @return 区间的值
     */
    public E query(int queryL, int queryR) {
        if (queryL < 0 || queryL > data.length || queryR < 0 || queryR > data.length) {
            throw new IllegalArgumentException("Query is illegal");
        }
        return query(0, 0, data.length - 1, queryL, queryR);
    }

    /**
     * 在以treeIndex为根的线段树中[l,r]的范围里，搜索区间[queryL, queryR]的值
     * @param treeIndex 根
     * @param l 左边界
     * @param r 右边界
     * @param queryL 要搜索区间的左边界
     * @param queryR 要搜索区间的右边界
     * @return 区间值
     */
    private E query(int treeIndex, int l, int r, int queryL, int queryR) {
        // 中止条件：当根节点的左边界和要查找的左边界重合 并且 当根节点的右边界和要查找的右边界重合
        if (l == queryL && r == queryR) {
            return tree[treeIndex];
        }
        // 当前根节点的左右孩子分隔索引
        int mid = l + (r - l) / 2;
        // 左孩子索引
        int leftTreeIndex = leftChild(treeIndex);
        // 右孩子索引
        int rightTreeIndex = rightChild(treeIndex);

        if (queryL >= mid + 1) {
            // 要查找的区间左节点大于中间节点
            return query(rightTreeIndex, mid + 1, r, queryL, queryR);
        } else if (queryR <= mid) {
            // 要查找的区间右节点小于中间节点
            return query(leftTreeIndex, l, mid, queryL, queryR);
        }
        // 查询的区间横跨左右孩子节点
        // 查询左孩子要搜索的区间
        E leftResult = query(leftTreeIndex, l, mid, queryL, mid);
        // 查询右孩子要搜索的区间
        E rightResult = query(rightTreeIndex, mid + 1, r, queryR, mid);

        // 融合两个值
        return merger.merge(leftResult, rightResult);
    }

    /**
     * 将index位置的值，更新为e
     * @param index 索引
     * @param e 值
     */
    public void set(int index, E e) {
        if (index < 0 || index >= data.length) {
            throw new IllegalArgumentException("Index is illegal");
        }
        data[index] = e;
    }

    /**
     * 在以treeIndex为根的线段树中更新index的值为e
     * @param treeIndex 根的索引
     * @param l 区间左边界
     * @param r 区间右边界
     * @param index 待更新位置的索引
     * @param e 待更新的值
     */
    private void set(int treeIndex, int l, int r, int index, E e) {
        // 最小问题：左右边界相等，说明只有一个元素
        if (l == r) {
            tree[treeIndex] = e;
            return;
        }

        // 孩子节点的分隔位置
        int mid = l + (r -l) / 2;
        // 左孩子节点位置
        int leftTreeIndex = leftChild(treeIndex);
        // 右孩子节点位置
        int rightTreeIndex = rightChild(treeIndex);
        if (index >= mid + 1) {
            set(rightTreeIndex, mid + 1, r, index, e);
        } else {
            // index <= mid
            set(leftTreeIndex, l, mid, index, e);
        }

        // 由于孩子节点发生了变化，那么对应父节点的位置同样也要发生变化
        tree[treeIndex] = merger.merge(tree[leftTreeIndex], tree[rightTreeIndex]);
    }

    @Override
    public String toString() {
        StringBuilder res = new StringBuilder();
        generateSegmentTreeString(0, 0, data.length - 1, 0, res);
        return res.toString();
    }

    private void generateSegmentTreeString(int treeIndex, int l, int r, int depth, StringBuilder res) {
        if (l == r) {
            res.append(generateDepthString(depth)).append(tree[treeIndex]).append("\n");
            return;
        }

        int leftTreeIndex = leftChild(treeIndex);
        int rightTreeIndex = rightChild(treeIndex);
        int mid = l + (r - l) / 2;

        res.append(generateDepthString(depth)).append("[").append(l).append(", ").append(r).append("] = ").append(tree[treeIndex]).append("\n");
        generateSegmentTreeString(leftTreeIndex, l, mid, depth + 1, res);
        generateSegmentTreeString(rightTreeIndex, mid + 1, r, depth + 1, res);
    }

    private String generateDepthString(int depth) {
        StringBuilder res = new StringBuilder();
        for (int i = 0; i < depth; i++) {
            res.append("--");
        }
        return res.toString();
    }
}
